/*
* cholesky_fact.cpp
*
*  Created on: Jun 22, 2013
*      Author: tifuera
*/

#include <iostream>
#include "sparse_band.h"
#include "mkl.h"
#include "cublas_v2.h"
#include "cuda_runtime.h"
#include "gpu_util.h"

void print_matrix(int rows, int cols, double * matrix, int ld)
{
	printf("PRINT START\n");
	for(int i = 0; i < rows; i++){
		for(int j = 0; j < cols; j++)
		{
			printf(" %6.1f",matrix[i*ld + j]);
		}
		printf("\n");
	}
	printf("PRINT END\n");
}


gpu_symm_band_matrix m2;//TODO remove
void print_tile(int tile_num)
{
	if(m2.order > 20)
	{
		return;
	}

	int num_tiles = (m2.order + m2.block_size - 1) / m2.block_size;
	int cur_bs = m2.block_size;
	if(tile_num == num_tiles -1)
	{
		cur_bs = m2.order - (num_tiles -1)*m2.block_size; 
	}

	double* tileCopy = (double *)malloc(m2.block_size*m2.tile_len[tile_num]*sizeof(double));
	//TODO copy it like a matrix
	for (int j = 0; j < m2.tile_len[tile_num]; j++)
	{
		cublasGetVector(cur_bs, sizeof(double),
			&m2.gpu_matrix_tiles[tile_num][j * cur_bs], 1,
			&tileCopy[j * m2.block_size], 1);
	}
	print_matrix(m2.tile_len[tile_num], cur_bs, tileCopy, m2.block_size);
}

void trsm_diagonal_block(double * h_block, double* h_last_block, int order,
	int lb_order, int ld)
{
	double one = 1.;
	dtrsm("L", "U", "T", "N", &order, &lb_order, &one, h_block, &ld, h_last_block, &ld);
}

void factorize_diagonal_block(double * h_block, int order, int ld)
{
	int mkl_info;
	dpotrf("U", &order, h_block, &ld, &mkl_info);
	if (mkl_info)
	{
		printf("BLOCK FACTORIZATION FAILED %d", mkl_info);
	}
}

void factorize(gpu_symm_band_matrix gpu_matrix, cublasHandle_t handle)
{
	cublasStatus_t status;
	double *h_block;
	double *h_last_block;
	int bs = gpu_matrix.block_size;
	int order = gpu_matrix.order;
	int hb = gpu_matrix.half_bandwith;
	int i, j;
	int num_tiles;
	int cur_bs;
	int cur_row;
	int index;
	int l21_row_num;
	int l22_tile_row_num;
	int l22_tile_col_num;
	double alpha_1 = 1;
	double alpha_m1 = -1;
	int tile_num;
	int last_block_size;
	m2 = gpu_matrix;
	print_tile(3);//TODO REMOVE
	/*h_block =*/cudaMallocHost(&h_block, bs * bs * sizeof(double));
	/*h_last_block =*/cudaMallocHost(&h_last_block, bs * bs * sizeof(double));
	/*for (i = 0; i < bs; i++)
	{
	for (j = 0; j < bs; j++)
	{
	h_block[i * bs + j] = 0;
	}
	}*/

	num_tiles = (order + bs - 1) / bs;
	last_block_size = order - (num_tiles -1)*bs; 

	for (i = 0, cur_row = 0; i < num_tiles; i++, cur_row += bs)
	{
		//copy current diagonal block
		cur_bs = bs;
		if (cur_row + cur_bs > order)
		{
			cur_bs = order - cur_row;
		}

		status = cublasGetMatrix(cur_bs, cur_bs, sizeof(double), gpu_matrix.gpu_matrix_tiles[i], cur_bs, h_block, bs);
		if(status != CUBLAS_STATUS_SUCCESS)
		{
			printf("Copy diagonal block from device failed");
		}
		factorize_diagonal_block(h_block, cur_bs, bs);
		//copy back
		/*for (j = 0; j < cur_bs; j++)
		{
		cublasSetVector(j + 1, sizeof(double), &h_block[j * bs], 1,
		&gpu_matrix.gpu_matrix_tiles[i][j * cur_bs], 1);
		}*/
		status = cublasSetMatrix(cur_bs, cur_bs, sizeof(double), h_block, bs, gpu_matrix.gpu_matrix_tiles[i], cur_bs );
		if(status != CUBLAS_STATUS_SUCCESS)
		{
			printf("Copy diagonal block to device failed");
		}

		/*l21_row_num = min(hb + 1, gpu_matrix.tile_len[i]) - cur_bs;*/
		l21_row_num = gpu_matrix.tile_len[i] - cur_bs;
		//update block in tile
		cublasDtrsm_v2(handle, CUBLAS_SIDE_LEFT, CUBLAS_FILL_MODE_UPPER,
			CUBLAS_OP_T, CUBLAS_DIAG_NON_UNIT, cur_bs, l21_row_num,
			&alpha_1, &gpu_matrix.gpu_matrix_tiles[i][0], cur_bs,
			&gpu_matrix.gpu_matrix_tiles[i][cur_bs * cur_bs], cur_bs);

		//update A22 start
		j = cur_bs;
		tile_num = i + 1;
		while (j < gpu_matrix.tile_len[i])
		{
			//TODO implement
			l22_tile_row_num = bs;
			if (l22_tile_row_num > gpu_matrix.tile_len[tile_num])
			{
				l22_tile_row_num = gpu_matrix.tile_len[tile_num];
			}
			int ldc = cur_bs;
			if(tile_num == num_tiles -1)
			{
				ldc = last_block_size;

			}
			cublasDgemm_v2(handle, CUBLAS_OP_T, CUBLAS_OP_N, l22_tile_row_num,
				gpu_matrix.tile_len[i] - j, cur_bs, &alpha_m1,
				&gpu_matrix.gpu_matrix_tiles[i][j * cur_bs], cur_bs,
				&gpu_matrix.gpu_matrix_tiles[i][j * cur_bs], cur_bs,
				&alpha_1, gpu_matrix.gpu_matrix_tiles[tile_num],
				ldc);

			j += cur_bs;
			tile_num++;
		}
		//update A22 end
	}
}
